import asyncio
from llama_index.core.agent.workflow import FunctionAgent

# Context可以用来管理对话上下文
from llama_index.core.workflow import Context

from llama_index.llms.openai import OpenAI
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

# 读取data目录下面的不同类型的文件，读到文件中的内容，把内容和元数据放到documents列表中
documents = SimpleDirectoryReader("data").load_data()
# 把文档进行切片，切成一段段小文本片段，把文本片转成向量，把文本本身和转成的向量
index = VectorStoreIndex.from_documents(documents)
# 实现查询引擎 就是可接收用户输入的查询，然后把查询转成向量，和上面index的向量进行相似度匹配获取匹配的小文本片段，
# 并装成context生成提示词发给大模型提问获取结果
query_engine = index.as_query_engine()


def multiply(a: float, b: float) -> float:
    """用于计算两个数字的乘积"""
    print("a, b", a, b)
    return a * b


async def search_documents(query: str) -> str:
    """用于根据自然语言来检索文档内容并返回答案"""
    print("query", query)
    response = await query_engine.aquery(query)
    return str(response)


# 创建一个带有计算功能的智能代理
agent = FunctionAgent(
    tools=[multiply, search_documents],  # 指定代理可用的工具
    llm=OpenAI(model="gpt-4o-mini"),  # 指定使用大模型工具
    system_prompt="你是一个乐于助人的助手，可能进行数学计算，也可以用检索文档来回答用户问题",  # 定义系统提示，描述代理的行为
)


# 定义异步函数
async def main():
    response = await agent.run("红楼梦的主要人物有哪些?另外8乘以9等于多少?")
    print(str(response))


if __name__ == "__main__":
    asyncio.run(main())
